package com.chen.zhipu.controller;

import com.chen.zhipu.dto.Message;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.Generation;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.zhipuai.ZhipuAiChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.io.PrintWriter;
import java.util.List;
import java.util.Map;

@RestController
@CrossOrigin
public class ChatController {
 
    private final ZhipuAiChatClient chatClient;
 
    @Autowired
    public ChatController(ZhipuAiChatClient chatClient) {
        this.chatClient = chatClient;
    }
 
    @GetMapping(value = "/v1/generate", produces = "application/json")
    public Map generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        String call = chatClient.call(message);
 
 
        System.out.println(call);
        return Map.of("generation", call);
    }
 
    @GetMapping("/v1/prompt")
    public List<Generation> prompt(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        PromptTemplate promptTemplate = new PromptTemplate("Tell me a {adjective} joke about {topic}");
        Prompt prompt = promptTemplate.create(Map.of("adjective", "funny", "topic", "cats"));
        return chatClient.call(prompt).getResults();
    }
 
    @GetMapping("/v1/chat/completions")
    public Flux<ChatResponse> chatCompletions(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return chatClient.stream(prompt);
    }
 
    @GetMapping("/v1/chat/completionsTow")
    public SseEmitter chatCompletionsTow(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        SseEmitter emitter = new SseEmitter();
        Prompt prompt = new Prompt(new UserMessage(message));
        // 订阅Flux并处理数据
        chatClient.stream(prompt).subscribe(
                data -> {
                    try {
                        emitter.send(data);
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete
        );
        return emitter;
    }
 
 
    @PostMapping("/v1/chat/completionsPost")
    public Flux<ChatResponse> completionsPost(@RequestBody Message msg) {
        Prompt prompt = new Prompt(new UserMessage(msg.getMsg()));
 
        return chatClient.stream(prompt);
    }
 
 
    @PostMapping(value = "/completionsPost", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public void completionsPost(@RequestBody Message msg, HttpServletResponse response) throws IOException {
        Prompt prompt = new Prompt(new UserMessage(msg.getMsg()));
        // 获取 SSE 的 flux 流
        Flux<ChatResponse> flux = chatClient.stream(prompt);
 
        // 设置必要的 SSE 头部
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");
        response.setHeader("Content-Type", "text/event-stream");
 
        // 获取响应输出流
        PrintWriter writer = response.getWriter();
        flux.doOnNext(chatResponse -> {
                    try {
                        // 将 ChatResponse 转换为 SSE 格式并写入响应流
                        Generation result = chatResponse.getResult();
                        AssistantMessage output = result.getOutput();
                        String content = output.getContent();
                        System.out.println("content: " + content);
                        writer.write("data: " + content + "\n\n");
                        writer.flush();
                    } catch (Exception e) {
                        // 处理异常
                        e.printStackTrace();
                    }
                })
                .doOnError(throwable -> {
                    try {
                        // 在出现错误时关闭响应流
                        writer.write("event: error\ndata: " + throwable.getMessage() + "\n\n");
                        writer.flush();
                        writer.close();
                    } catch (Exception e) {
                        // 处理异常
                        e.printStackTrace();
                    }
                })
                .doOnComplete(() -> {
                    try {
                        // 在流完成时关闭响应流
                        writer.write("event: complete\ndata: Stream completed\n\n");
                        writer.flush();
                        writer.close();
                    } catch (Exception e) {
                        // 处理异常
                        e.printStackTrace();
                    }
                })
                .subscribe(); // 订阅以启动流
    }
 
 
    @PostMapping(value = "/completionsPostNew", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter completionsPostNew(@RequestBody Message msg, HttpServletResponse response) {
        SseEmitter emitter = new SseEmitter();
        UserMessage userMessage = new UserMessage(msg.getMsg());
        SystemMessage systemMessage = new SystemMessage("请用中文回答。你是一位智能编程助手，你叫CSBlogAi。你会为用户回答关于编程、代码、计算机方面的任何问题，" +
                "并提供格式规范、可以执行、准确安全的代码，并在必要时提供详细的解释。" +
                "任务：请为输入代码提供格式规范的注释，包含多行注释和单行注释，" +
                "请注意不要改动原始代码，只需要添加注释。 ");
        List<org.springframework.ai.chat.messages.Message> messages = List.of(userMessage, systemMessage);
        Prompt prompt = new Prompt(messages);
        // 设置响应头
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        // 订阅Flux并处理数据
        chatClient.stream(prompt).subscribe(
                data -> {
                    try {
                        // 发送 JSON 格式的 SSE 事件
                        emitter.send(data);
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete
        );
        return emitter;
    }


    @PostMapping(value = "/completionsPostShop", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter completionsPostShop(@RequestBody Message msg, HttpServletResponse response) {
        SseEmitter emitter = new SseEmitter();
        UserMessage userMessage = new UserMessage(msg.getMsg());
        SystemMessage systemMessage = new SystemMessage("请用中文回答。你是一位商城助手，主要负责商品的介绍和解释对应的做法，生源地等。多扩展商品明细，详细介绍商品~");
        List<org.springframework.ai.chat.messages.Message> messages = List.of(userMessage, systemMessage);
        Prompt prompt = new Prompt(messages);
        // 设置响应头
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        // 订阅Flux并处理数据
        chatClient.stream(prompt).subscribe(
                data -> {
                    try {
                        // 发送 JSON 格式的 SSE 事件
                        emitter.send(data);
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete
        );
        return emitter;
    }
 
}